package com.cyz.manager.sharding;

import com.baomidou.mybatisplus.extension.toolkit.SqlRunner;
import com.cyz.exception.BusinessException;
import com.cyz.exception.ErrorCode;
import com.cyz.model.entity.Space;
import com.cyz.model.enums.SpaceLevelEnum;
import com.cyz.model.enums.SpaceTypeEnum;
import com.cyz.service.SpaceService;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.driver.jdbc.core.connection.ShardingSphereConnection;
import org.apache.shardingsphere.infra.metadata.database.rule.ShardingSphereRuleMetaData;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.rule.ShardingRule;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.Collections;
import java.util.List;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 自定义动态分表管理器
 * 1. 将管理器注册为 Bean，通过 @PostConstruct 注解，在 Bean 加载后获取所有的分表并更新配置。
 * 2. 编写获取分页列表的方法， 从数据库中查询符合要求的空间列表，再拼接上逻辑表，就得到完整的分页列表。
 * 3. 更新 ShardingSphere 的 actual-data-nodes 动态表名配置。获取到 ShardingSphere 的ContextManager，
 * 找到配置文件中的 actual-data-nodes 进行更新。从而实现自定义分库分表管理。
 */
@Slf4j
//@Component   // 为了方便部署，暂时不使用分库分表
public class DynamicShardingManager {

    @Resource
    private DataSource dataSource;

    @Resource
    private SpaceService spaceService;

    private static final String LOGIC_TABLE_NAME = "picture";

    private static final String DATABASE_NAME = "logic_db"; // 代表配置文件中的数据库名称

    /**
     * 初始化项目的时候，动态设置分表配置
     */
    @PostConstruct
    public void initialize(){
        log.info("初始化项目的时候，动态分表配置");
        updateShardingTableNodes();
    }

    /**
     * 获取所有需要分的动态表名（ picture 和 picture_${空间ID} ）
     */
    private Set<String> fetchAllPictureTableName(){
        // 查询数据库，获取表名
        Set<Long> spaceIds = spaceService.lambdaQuery()
                .eq(Space::getSpaceType, SpaceTypeEnum.TEAM.getValue())
                .list().stream()
                .map(Space::getId)
                .collect(Collectors.toSet());

        // 拼接表名
        Set<String> spaceTableNames = spaceIds.stream()
                .map(spaceId -> LOGIC_TABLE_NAME + "_" + spaceId)
                .collect(Collectors.toSet());
        // 添加初始逻辑表名
        spaceTableNames.add(LOGIC_TABLE_NAME);
        return spaceTableNames;
    }

    /**
     * 动态设置分表配置方法   大部分为模板代码
     */
    private void updateShardingTableNodes() {
        // 获取所有动态表名
        Set<String> spaceTableNames = fetchAllPictureTableName();
        // 拼接成 配置中的值
        String newActualDataNodes = spaceTableNames.stream()
                .map(tableName -> "yun_picture." + tableName)   // 必须加 . 确保前缀合法
                .collect(Collectors.joining(","));

        log.info("动态分表 actual-data-nodes 配置：{}", newActualDataNodes);


        // 以下为模板代码，动态修改分库分表的配置 actual-data-nodes 属性
        // 获取 ShardingSphere ContextManager
        ContextManager contextManager = getContextManager();
        ShardingSphereRuleMetaData ruleMetaData = contextManager.getMetaDataContexts()
                .getMetaData()
                .getDatabases()
                .get(DATABASE_NAME)
                .getRuleMetaData();

        Optional<ShardingRule> shardingRule = ruleMetaData.findSingleRule(ShardingRule.class);
        if (shardingRule.isPresent()) {
            ShardingRuleConfiguration ruleConfig = (ShardingRuleConfiguration) shardingRule.get().getConfiguration();
            List<ShardingTableRuleConfiguration> updatedRules = ruleConfig.getTables()
                    .stream()
                    .map(oldTableRule -> {
                        if (LOGIC_TABLE_NAME.equals(oldTableRule.getLogicTable())) {
                            ShardingTableRuleConfiguration newTableRuleConfig = new ShardingTableRuleConfiguration(LOGIC_TABLE_NAME, newActualDataNodes);
                            newTableRuleConfig.setDatabaseShardingStrategy(oldTableRule.getDatabaseShardingStrategy());
                            newTableRuleConfig.setTableShardingStrategy(oldTableRule.getTableShardingStrategy());
                            newTableRuleConfig.setKeyGenerateStrategy(oldTableRule.getKeyGenerateStrategy());
                            newTableRuleConfig.setAuditStrategy(oldTableRule.getAuditStrategy());
                            return newTableRuleConfig;
                        }
                        return oldTableRule;
                    })
                    .collect(Collectors.toList());
            ruleConfig.setTables(updatedRules);
            contextManager.alterRuleConfiguration(DATABASE_NAME, Collections.singleton(ruleConfig));
            contextManager.reloadDatabase(DATABASE_NAME);
            log.info("动态分表规则更新成功！");
        } else {
            log.error("未找到 ShardingSphere 的分片规则配置，动态分表更新失败。");
        }
    }

    /**
     * 获取 ShardingSphere ContextManager
     */
    private ContextManager getContextManager(){
        try(ShardingSphereConnection connection = dataSource.getConnection().unwrap(ShardingSphereConnection.class)){
            return connection.getContextManager();
        }catch (SQLException e){
            throw new RuntimeException("获取 ShardingSphere ContextManager 失败", e);
        }
    }

    /**
     * 动态建表，为团队空间动态创建图片分表
     */
    public void createSpacePictureTable(Space space){
        // 判断是否是旗舰版的团队空间，仅为旗舰空间开通独立空间
        if (space.getSpaceType() == SpaceTypeEnum.TEAM.getValue() && space.getSpaceLevel() == SpaceLevelEnum.FLAGSHIP.getValue()){
            Long spaceId = space.getId();
            // 拼接表名
            String tableName = "picture_" + spaceId;
            // 创建新的团队空间图片表
            String createTableSQL = "CREATE TABLE " + tableName + " LIKE picture";
            try{
                log.info("创建分表：{}", createTableSQL);
                SqlRunner.db().update(createTableSQL);
                // 更新分表
                updateShardingTableNodes();
            }catch (Exception e){
                log.error("创建团队空间表失败", e);
                log.error("创建团队空间表失败，空间 id = {}", spaceId);
                throw new BusinessException(ErrorCode.OPERATION_ERROR, "团队空间图片表创建失败");
            }
        }
    }

}
